perf(refit): batch MXFP8 expert prequantization - #3804

Draft

seonjinn wants to merge 104 commits into

NVIDIA-NeMo:mainfrom

seonjinn:sna/perf-mxfp8-batched-expert-prequantization

seonjinn commented Aug 24, 2026 •

edited

Loading

Contributor

Summary

Batch matching MoE expert weights before trainer-side MXFP8 prequantization. The path quantizes at most 16 experts per call, preserves the original refit wire order, and reuses temporary stacking storage only within one export pass.

Non-expert weights and expert groups that cannot be batched keep the existing per-tensor path. fp8_param=true is rejected because this path requires BF16 source weights.

This PR builds on #3294. Async NCCL Reshard receiver-side conversion is separate and is not changed here.

Why

The existing path launches one quantization operation per expert tensor. Large MoE models repeat this operation hundreds of times per refit. Batching compatible experts reduces launch and synchronization overhead without changing the value/scale payload consumed by vLLM.

Evidence

Prior Qwen3-30B-A3B prototype on GB200, synchronous colocated GRPO with CUDA Graph, 20 steps:

Metric	Per expert	Batched	Change
Transfer/update	4.506 s	2.200 s	-51.2%
Total refit	8.431 s	5.935 s	-29.6%
E2E step	188.88 s	185.32 s	-1.89%
Throughput	2,199.0 tok/s/GPU	2,243.0 tok/s/GPU	+2.00%

Generation KL error was unchanged at 0.00398.

Current head cfbe773c70a61f55835e6ed15339c61e3ae844ed:

48 relevant GB200 unit tests passed, including quantization parity, Bridge export order, bounded batching, split payloads, scale shape, CUDA stream handoffs, and recipe scope.
Ruff lint, format, and diff checks passed on all changed files.
Current-head 20-step performance A/B is running; the table above remains prototype evidence until that rerun completes.

Qwen3-235B peak-memory and end-to-end validation remain before this PR is ready for review.

seonjinn added 30 commits

July 21, 2026 16:16


          perf(vllm): optimize quantized refit paths

b88aa2c

Reduce MXFP8 and ModelOpt refit overhead while preserving transport and checkpoint-engine lifecycle correctness.

Signed-off-by: sna <sna@nvidia.com>


          Merge branch 'main' into sna/pr-mxfp8-refit-optimization

438e20d


          Merge branch 'main' into sna/pr-mxfp8-refit-optimization

96ba299

Signed-off-by: sna <sna@nvidia.com>


          fix(trtllm): align refit metadata interface

0fb59b6

Signed-off-by: sna <sna@nvidia.com>


          test(vllm): include MXFP8 refit tests in L0

c4a5e14

Signed-off-by: sna <sna@nvidia.com>


          fix(refit): complete prequant metadata handshake

5c53bd4

Signed-off-by: sna <sna@nvidia.com>


          test(refit): cover MXFP8 optimization paths

98fc197

Signed-off-by: sna <sna@nvidia.com>


          test(refit): fix vLLM test import order

b237fd2

Signed-off-by: sna <sna@nvidia.com>


          Merge upstream/main into sna/pr-mxfp8-refit-optimization

d1f9fc9

Signed-off-by: sna <sna@nvidia.com>


          test(fp8): cover MXFP8 MoE padding path

ab030c3

Signed-off-by: sna <sna@nvidia.com>


          Merge upstream/main into sna/pr-mxfp8-refit-optimization

2c147ac

Signed-off-by: sna <sna@nvidia.com>


          Merge remote-tracking branch 'upstream/main' into sna/pr-mxfp8-refit-…

150e25f

…optimization-squashed

# Conflicts:
#	nemo_rl/algorithms/grpo.py


          fix(refit): reconcile MXFP8 optimizations with NCCL reshard

f391a4a

Signed-off-by: seonjinn <sna@nvidia.com>


          fix(refit): validate checkpoint prequant backend

85a6fea

Signed-off-by: seonjinn <sna@nvidia.com>


          fix(refit): reject MXFP8 with NCCL reshard

716bb74

Signed-off-by: seonjinn <sna@nvidia.com>


          fix(refit): validate NCCL reshard storage precision

f638fa6

Signed-off-by: seonjinn <sna@nvidia.com>


          ci: associate Codecov uploads with pull requests

2fde3b0

Signed-off-by: seonjinn <sna@nvidia.com>


          Merge remote-tracking branch 'origin/main' into sna/pr-mxfp8-refit-op…

2b0b262

…timization-squashed

Signed-off-by: seonjinn <sna@nvidia.com>


          Merge remote-tracking branch 'upstream/main' into sna/pr-mxfp8-refit-…

6fe0632

…optimization-squashed

Signed-off-by: seonjinn <sna@nvidia.com>


          fix: address MXFP8 refit review feedback

00fa128

Signed-off-by: seonjinn <sna@nvidia.com>


          Merge remote-tracking branch 'upstream/main' into sna/pr-mxfp8-refit-…

de65f06

…optimization-squashed

Signed-off-by: seonjinn <sna@nvidia.com>


          Merge remote-tracking branch 'upstream/main' into sna/pr-mxfp8-refit-…

…optimization-squashed

Signed-off-by: seonjinn <sna@nvidia.com>

# Conflicts:
#	nemo_rl/models/generation/vllm/quantization/fp8.py
#	tests/unit/models/generation/test_vllm_fp8_quantization.py


          fix(vllm): preserve MXFP8 refit on vLLM 0.25

Signed-off-by: seonjinn <sna@nvidia.com>


          style(vllm): sort MXFP8 backend import

123cc07

Signed-off-by: seonjinn <sna@nvidia.com>


          fix(vllm): allow partial configs in quant validation

ef3fa08

Signed-off-by: seonjinn <sna@nvidia.com>


          Merge remote-tracking branch 'upstream/main' into sna/pr-mxfp8-refit-…

…optimization-squashed


          feat(recipe): add async Qwen3 30B MXFP8 rollout

c4a7b06

Signed-off-by: seonjinn <sna@nvidia.com>


          feat(recipe): add async Qwen3 235B MXFP8 rollout

69f7995

Signed-off-by: seonjinn <sna@nvidia.com>


          ci: allowlist async recipe result path

d53bfdf

Signed-off-by: seonjinn <sna@nvidia.com>


          fix(vllm): initialize async driver FP8 config

2405bf5

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn commented Aug 29, 2026

Contributor Author

/ok to test 1c9e4b5

seonjinn added 27 commits

September 4, 2026 04:55


          Merge remote-tracking branch 'origin/main' into codex/pr3294-push-202…

7c9b881

…60904

Signed-off-by: seonjinn <sna@nvidia.com>

# Conflicts:
#	nemo_rl/algorithms/grpo.py
#	nemo_rl/weight_sync/collective_weight_synchronizer.py
#	nemo_rl/weight_sync/interfaces.py
#	tests/unit/models/generation/test_vllm_fp8_quantization.py


          fix(mxfp8): handle mixed grouped expert refit

bf03429

Signed-off-by: seonjinn <sna@nvidia.com>


          test(mxfp8): align refit fixtures with vllm

d9246a9

Signed-off-by: seonjinn <sna@nvidia.com>


          test(mxfp8): model routed expert config

443a9f9

Signed-off-by: seonjinn <sna@nvidia.com>


          fix(vllm): reroute grouped MXFP8 scale sidecars

b8159de

Signed-off-by: seonjinn <sna@nvidia.com>


          test(mxfp8): patch modular kernel factory

8d38b28

Signed-off-by: seonjinn <sna@nvidia.com>


          test(mxfp8): use block-aligned hidden size

658d385

Signed-off-by: seonjinn <sna@nvidia.com>


          test(mxfp8): exercise unpadded kernel reuse

172d346

Signed-off-by: seonjinn <sna@nvidia.com>


          test(mxfp8): use current refit loader contract

546d346

Signed-off-by: seonjinn <sna@nvidia.com>


          test(mxfp8): update remaining loader fixtures

db15416

Signed-off-by: seonjinn <sna@nvidia.com>


          Merge remote-tracking branch 'upstream/main' into sna/pr3804-replay-2…

b8247a6

…0260905


          perf(refit): batch MXFP8 expert prequantization

d1c11a2

Signed-off-by: seonjinn <sna@nvidia.com>


          fix(refit): support grad-enabled expert views

3b4c992

Signed-off-by: seonjinn <sna@nvidia.com>


          fix(refit): release MXFP8 scratch after export

3c7015f

Signed-off-by: seonjinn <sna@nvidia.com>


          fix(refit): preserve expert export and scale contracts

510cc4b

Signed-off-by: seonjinn <sna@nvidia.com>


          fix(refit): synchronize batched outputs across streams

5b2036c

Signed-off-by: seonjinn <sna@nvidia.com>


          fix(refit): synchronize fallback and pending tensors

6f6a58a

Signed-off-by: seonjinn <sna@nvidia.com>


          test(refit): use BF16 grad-enabled expert fixtures

72f44ea

Signed-off-by: seonjinn <sna@nvidia.com>


          test(refit): cover split MXFP8 weight and scale batches

b579505

Signed-off-by: seonjinn <sna@nvidia.com>


          fix(refit): allow split MXFP8 weight scale batches

230f275

Signed-off-by: seonjinn <sna@nvidia.com>


          test(refit): reject MXFP8 parameter prequantization

0dd1c01

Signed-off-by: seonjinn <sna@nvidia.com>


          test(refit): cover batched prequantization contracts

f4cc305

Signed-off-by: seonjinn <sna@nvidia.com>


          fix(refit): preserve batched prequant wire order

f0d40b0

Signed-off-by: seonjinn <sna@nvidia.com>


          style: format refit prequantization changes

ebdbfc3

Signed-off-by: seonjinn <sna@nvidia.com>


          test(refit): follow preserved expert wire order

887a00b

Signed-off-by: seonjinn <sna@nvidia.com>


          fix(refit): honor disabled FP8 parameter config

eecb63a

Signed-off-by: seonjinn <sna@nvidia.com>


          Merge latest refit base into batched prequantization

cfbe773

Use the replayed PR NVIDIA-NeMo#3804 tree on the latest PR NVIDIA-NeMo#3294 and main integration.

Signed-off-by: seonjinn <sna@nvidia.com>

github-actions Bot removed the CI label

seonjinn commented Sep 5, 2026

Contributor Author

Self-review completed on head cfbe773c70a61f55835e6ed15339c61e3ae844ed.

Reviewed the six-file change on top of #3294, including numerical behavior, refit wire order, CUDA stream handoffs, temporary-buffer lifetime, configuration guards, and fallback paths. No blocking correctness, compatibility, security, or deadlock issue was found.

Key checks:

The batched iterator preserves the exact Bridge export order. This is required because the collective sender transfers tensors without parameter names.
Each quantization call is bounded to 16 compatible experts; incompatible or non-expert tensors keep the per-tensor path.
Split value/scale batches and singleton scale dimensions retain the existing vLLM payload contract.
Active fp8_param=true is rejected because trainer-side prequantization requires BF16 source weights. A stale fp8_param=true value is ignored when FP8 is disabled.
Producer/consumer CUDA stream edges and temporary scratch lifetime have direct regression coverage.

Validation on GB200:

48 passed, 29 deselected in the relevant vLLM/MXFP8 unit-test selection.
The focused source-order, bounded-batch, scale-parity, and stream-handoff selection passed 9/9.
Ruff lint, format, and git diff --check passed for all changed files.
All commits are signed off; GitHub DCO, title, copyright, secrets, and submodule checks pass.

The current-head Qwen3-30B-A3B 20-step performance A/B and Qwen3-235B peak-memory validation are still pending, so the PR remains a draft.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:L0 Documentation Low Precision Performance